package com.magenta.maxoptra.gis.distancematrix.iterators;

import com.magenta.maxoptra.gis.commons.Pair;
import com.magenta.maxoptra.gis.distancematrix.DistanceMatrixCoordinate;
import com.magenta.maxoptra.gis.distancematrix.RoutesIterator;
import com.magenta.maxoptra.gis.distancematrix.StateAwareRouteIterator;
import org.jetbrains.annotations.Nullable;

import java.util.*;

/**
 * @author Sergey Grachev
 */
public final class AsymmetricIterator implements RoutesIterator, StateAwareRouteIterator {

    private final List<DistanceMatrixCoordinate> forwardCoordinates;
    private final List<DistanceMatrixCoordinate> to;
    private final List<DistanceMatrixCoordinate> currentTo;
    private final List<DistanceMatrixCoordinate> backwardCoordinates;
    private final Map<Type, List<DistanceMatrixCoordinate>> finished = new HashMap<Type, List<DistanceMatrixCoordinate>>();
    private final List<DistanceMatrixCoordinate> finishedForward = new ArrayList<DistanceMatrixCoordinate>(1);
    private final List<DistanceMatrixCoordinate> finishedBackward = new ArrayList<DistanceMatrixCoordinate>(1);

    {
        finished.put(Type.FORWARD, finishedForward);
        finished.put(Type.BACKWARD, finishedBackward);
    }

    private final Map<Type, List<DistanceMatrixCoordinate>> finishedUnmodifiable = Collections.unmodifiableMap(finished);
    private final long totalSize;

    private long totalIndex;
    private InternalState savedState;
    private Type type = Type.FORWARD;
    private int fromIndex = 0;
    private int toIndex = -1;

    public AsymmetricIterator(final List<DistanceMatrixCoordinate> from, final List<DistanceMatrixCoordinate> to) {
        this.to = to;
        this.forwardCoordinates = new ArrayList<DistanceMatrixCoordinate>(from);
        this.backwardCoordinates = new ArrayList<DistanceMatrixCoordinate>(to);
        this.currentTo = new ArrayList<DistanceMatrixCoordinate>(to.size());
        this.totalSize = calculateSize();

        moveToNext(true);
    }

    private long calculateSize() {
        long result = 0;
        long size;

        // forward

        size = 0;
        final Set<DistanceMatrixCoordinate> backward = new HashSet<DistanceMatrixCoordinate>(to);
        for (final DistanceMatrixCoordinate item : forwardCoordinates) {
            if (item.isNeedCalculateForward()) {
                backward.remove(item);
                size++;
            }
        }

        result += size * to.size() - (to.size() - backward.size());

        // backward

        size = 0;
        for (final DistanceMatrixCoordinate item : backward) {
            if (item.isNeedCalculateBackward()) {
                size++;
            }
        }

        return result + forwardCoordinates.size() * size;
    }

    private boolean moveToNext(final boolean initialize) {
        if (!hasNext()) {
            return false;
        }

        if (!initialize) {
            finishedForward.clear();
            finishedBackward.clear();
        }

        switch (type) {
            case FORWARD: {
                if (toIndex + 1 >= currentTo.size()) { // row is finished, move to next

                    if (!initialize) {
                        final DistanceMatrixCoordinate current = forwardCoordinates.get(fromIndex);
                        if (current.isNeedCalculateForward()) {
                            finishedForward.add(current);
                        }
                        backwardCoordinates.remove(current);
                        fromIndex++;
                    }

                    toIndex = 0;
                    currentTo.clear();

                    for (; fromIndex < forwardCoordinates.size(); fromIndex++) {
                        final DistanceMatrixCoordinate next = forwardCoordinates.get(fromIndex);
                        if (next.isNeedCalculateForward()) {
                            currentTo.addAll(to);
                            currentTo.remove(next);
                            if (currentTo.isEmpty()) {
                                continue;
                            }
                            return true;
                        }
                    }

                    // no more forward routes, try backward routes
                    type = Type.BACKWARD;
                    currentTo.clear();
                    fromIndex = 0;
                    toIndex = -1;

                    if (!hasNext()) {
                        addFinishedBackwardCoordinates();
                    }

                    moveToNext(true);

                    return true;
                }

                toIndex++;

                break;
            }

            case BACKWARD: {
                if (toIndex + 1 >= currentTo.size()) { // row is finished, move to next

                    if (!initialize) {
                        final DistanceMatrixCoordinate current = backwardCoordinates.get(fromIndex);
                        if (current.isNeedCalculateBackward()) {
                            finishedBackward.add(current);
                        }
                        fromIndex++;
                    }

                    toIndex = 0;
                    currentTo.clear();

                    for (; fromIndex < backwardCoordinates.size(); fromIndex++) {
                        final DistanceMatrixCoordinate next = backwardCoordinates.get(fromIndex);
                        if (next.isNeedCalculateBackward()) {
                            currentTo.addAll(forwardCoordinates);
                            currentTo.remove(next);
                            if (currentTo.isEmpty()) {
                                continue;
                            }
                            return true;
                        }
                    }

                    // skipped coordinates
                    addFinishedBackwardCoordinates();

                } else {
                    toIndex++;
                }
            }
        }

        return false;
    }

    private void addFinishedBackwardCoordinates() {
        final List<DistanceMatrixCoordinate> skipped = new ArrayList<DistanceMatrixCoordinate>(to);
        skipped.removeAll(backwardCoordinates);
        for (final DistanceMatrixCoordinate coordinate : skipped) {
            if (coordinate.isNeedCalculateBackward()) {
                finishedBackward.add(coordinate);
            }
        }
    }

    private void fillCoordinates(final int limit, final List<DistanceMatrixCoordinate> toCoordinates) {
        int i = 0;
        do {
            if (!hasNext()) {
                break;
            }
            totalIndex++;
            toCoordinates.add(currentTo.get(toIndex));
        } while (!moveToNext(false) && ++i < limit);
    }

    @Nullable
    @Override
    public Pair<DistanceMatrixCoordinate, DistanceMatrixCoordinate> nextRoute() {
        if (!hasNext()) {
            return null;
        }

        final DistanceMatrixCoordinate fromCoordinate, toCoordinate;

        switch (type) {
            case FORWARD: {
                fromCoordinate = forwardCoordinates.get(fromIndex);
                toCoordinate = currentTo.get(toIndex);
                break;
            }

            case BACKWARD: {
                fromCoordinate = backwardCoordinates.get(fromIndex);
                toCoordinate = currentTo.get(toIndex);
                break;
            }

            default:
                throw new IllegalStateException("Unknown state " + type);
        }

        moveToNext(false);

        totalIndex++;

        return new Pair<DistanceMatrixCoordinate, DistanceMatrixCoordinate>(fromCoordinate, toCoordinate);
    }

    @Nullable
    @Override
    public Pair<DistanceMatrixCoordinate, List<DistanceMatrixCoordinate>> nextRoutes(final int limit) {
        if (!hasNext()) {
            return null;
        }

        final DistanceMatrixCoordinate fromCoordinate;
        final List<DistanceMatrixCoordinate> toCoordinates = new ArrayList<DistanceMatrixCoordinate>();

        switch (type) {
            case FORWARD: {
                fromCoordinate = forwardCoordinates.get(fromIndex);
                fillCoordinates(limit, toCoordinates);
                break;
            }

            case BACKWARD: {
                fromCoordinate = backwardCoordinates.get(fromIndex);
                fillCoordinates(limit, toCoordinates);
                break;
            }

            default:
                throw new IllegalStateException("Unknown state " + type);
        }

        return new Pair<DistanceMatrixCoordinate, List<DistanceMatrixCoordinate>>(fromCoordinate, toCoordinates);
    }

    @Override
    public boolean hasNext() {
        return Type.FORWARD.equals(type)
                ? !forwardCoordinates.isEmpty() && (fromIndex < forwardCoordinates.size() || toIndex < currentTo.size())
                : !backwardCoordinates.isEmpty() && (fromIndex < backwardCoordinates.size() || toIndex < currentTo.size());
    }

    @Override
    public boolean hasFinishedCoordinates() {
        return !finishedForward.isEmpty() || !finishedBackward.isEmpty();
    }

    @Override
    public Map<Type, List<DistanceMatrixCoordinate>> getFinishedCoordinates() {
        return finishedUnmodifiable;
    }

    @Override
    public long index() {
        return totalIndex;
    }

    @Override
    public long size() {
        return totalSize;
    }

    @Override
    public void saveState() {
        savedState = new InternalState(currentTo, backwardCoordinates, type, fromIndex, toIndex,
                finishedForward, finishedBackward, totalIndex);
    }

    @Override
    public void restoreState() {
        if (savedState != null) {
            currentTo.clear();
            currentTo.addAll(savedState.currentTo);
            backwardCoordinates.clear();
            backwardCoordinates.addAll(backwardCoordinates);
            finishedForward.clear();
            finishedForward.addAll(savedState.finishedForward);
            finishedBackward.clear();
            finishedBackward.addAll(savedState.finishedForward);
            type = savedState.type;
            fromIndex = savedState.fromIndex;
            toIndex = savedState.toIndex;
            totalIndex = savedState.totalIndex;
            savedState = null;
        }
    }

    private final static class InternalState {
        public final List<DistanceMatrixCoordinate> currentTo;
        public final List<DistanceMatrixCoordinate> backwardCoordinates;
        public final List<DistanceMatrixCoordinate> finishedForward;
        public final List<DistanceMatrixCoordinate> finishedBackward;
        public final Type type;
        public final int fromIndex;
        public final int toIndex;
        public final long totalIndex;

        private InternalState(final List<DistanceMatrixCoordinate> currentTo, final List<DistanceMatrixCoordinate> backwardCoordinates,
                              final Type type, final int fromIndex, final int toIndex,
                              final List<DistanceMatrixCoordinate> finishedForward, final List<DistanceMatrixCoordinate> finishedBackward,
                              final long totalIndex) {
            this.finishedForward = finishedForward;
            this.finishedBackward = finishedBackward;
            this.totalIndex = totalIndex;
            this.currentTo = new ArrayList<DistanceMatrixCoordinate>(currentTo);
            this.backwardCoordinates = new ArrayList<DistanceMatrixCoordinate>(backwardCoordinates);
            this.type = type;
            this.fromIndex = fromIndex;
            this.toIndex = toIndex;
        }
    }
}
